# -*- coding: utf-8 -*-
# ----------------------------
# @Time    : 2022/5/29 3:27 PM
# @Author  : changqingai
# @FileName: 02-from_genrator.py
# ----------------------------

import numpy as np
import tensorflow as tf


# ********* 例子1 ******************
# 定义生成器
def count(stop):
    i = 0
    while i < stop:
        yield i
        i += 1


# 生成tf.生成器
ds_count = tf.data.Dataset.from_generator(count, args=[25], output_types=tf.int32, output_shapes=())
for batch in ds_count.shuffle(buffer_size=100, seed=1).repeat().batch(4).take(10):
    print("batch:", batch.numpy())


# ********* 例子2 ******************
def gen():
    ragged_tensor = tf.ragged.constant([[1, 2], [3]])
    yield 42, ragged_tensor


ds2 = tf.data.Dataset.from_generator(gen,
                                     output_signature=(tf.TensorSpec(shape=(),
                                                                     dtype=tf.int32),
                                                       tf.RaggedTensorSpec(shape=(2, None),
                                                                           dtype=tf.int32)))

for batch in ds2.repeat(1):
    print("batch:", batch[0].numpy(), batch[1].numpy())


# ********* 例子3 ******************
def gen_series(stop):
    i = 0
    while True:
        size = np.random.randint(1, 5)
        yield i, np.random.normal(size=(size,))
        i += 1
        if i > stop:
            break


# 普通生成器调用
for i, series in gen_series(5):
    print("i: ", i, series)


series_ds = tf.data.Dataset.from_generator(gen_series, args=[5],
                                           output_types=(tf.int32, tf.float32),
                                           output_shapes=((), (None,)))

for batch in series_ds.repeat(2):
    print("batch:", batch[0].numpy(), batch[1].numpy())
